import os
import shutil

from paddleocr import PaddleOCR

import config
from progress_seq_tool import load_last_processed_seq, save_total_number, load_total_number, save_processed_seq


class SensitiveKeywordsChecker:
    def __init__(self, input_path, output_path):
        self.input_path_ = input_path
        self.output_path_ = output_path

        # 准备结果数据表

        self.paddle_ocr_ = PaddleOCR(use_angle_cls=True, lang="ch")

    def count_jpg_files(self):
        total_jpg_files = 0
        for root, dirs, files in os.walk(self.input_path_):
            for file in files:
                file_path = os.path.join(root, file)
                if file_path.lower().endswith('.jpg'):
                    total_jpg_files += 1

        save_total_number(self.output_path_,total_jpg_files)
        return total_jpg_files

    def check_all_jpg_files(self):
        total_jpg_files = load_total_number(self.output_path_)
        if total_jpg_files == 0:
            total_jpg_files = self.count_jpg_files()

        logging_info = fr'一共有{total_jpg_files}张jpg图片！'
        config.print_console_and_write_log(logging_info)

        processed_jpg_files = load_last_processed_seq(self.output_path_)

        for root, dirs, files in os.walk(self.input_path_):
            for file in files:
                file_path = os.path.join(root, file)
                if file_path.lower().endswith('.jpg'):
                    processed_jpg_files += 1
                    try:
                        self.check_one_jpg_file(file_path)

                        # 计算新的文件路径，保持相同的目录结构
                        relative_path = os.path.relpath(file_path, self.input_path_)
                        new_file_path = os.path.join(self.output_path_, relative_path)

                        # 确保目标目录存在
                        os.makedirs(os.path.dirname(new_file_path), exist_ok=True)

                        # 将文件从原目录移动到新目录
                        shutil.move(file_path, new_file_path)

                        progress = (processed_jpg_files / total_jpg_files) * 100
                        print(f"已处理 {processed_jpg_files} / {total_jpg_files} 个.jpg文件，进度：{progress:.2f}%")

                        save_processed_seq(self.output_path_,processed_jpg_files)

                    except Exception as e:
                        config.print_console_and_write_log(f"处理文件 {file_path} 时发生错误: {e}")

        config.print_console_and_write_log(f"检查完成，共处理了 {total_jpg_files} 个.jpg文件")



    def check_one_jpg_file(self, image_file_path):

        """
        识别当前的图片中的文字，检查是否有敏感关键字。
        """

        try:
            result = self.paddle_ocr_.ocr(image_file_path, cls=True)
        except Exception as e:
            logging_info = f"OCR 过程中发生错误：{e}"
            self.print_console_and_write_log(logging_info)

            return

        sensitive_words = ['绝密', '机密', '秘密']
        for idx in range(len(result)):
            res = result[idx]
            for line in res:
                target = line[1][0]
                if isinstance(target, str):
                    for word in sensitive_words:
                        if word in target:
                            self.result_seq_ += 1
                            row = [self.result_seq_, image_file_path, word]
                            self.worksheet_.append(row)
                            logging_info = f'找到{word}资料一份，路径为：{image_file_path}'
                            config.print_console_and_write_log(logging_info)
                            return  # 找到敏感词后，可以退出内层循环
